import random
import sys
import math
import pygame
from pygame.color import *

# 常量定义
G = 8000  # 引力常数，为了模拟效果放大
WINDOW_SIZE = (1500, 850)  # 窗口大小
TIME_SLOT = 0.02  # 时间片，每次模拟的时间间隔
INITIAL_ACCURACY = 1  # 初始精度
MASS = 10 ** 3  # 恒星质量
RADIUS = 8  # 恒星半径

# 恒星类定义
class Star:
    ax, ay = 0, 0  # 初始x,y方向上的加速度

    def __init__(self, x, y, vx, vy, mass, radius):
        self.x = x  # x坐标
        self.y = y  # y坐标
        self.vx = vx  # x方向速度
        self.vy = vy  # y方向速度
        self.mass = mass  # 质量
        self.radius = radius  # 半径

    def acceleration(self, other_star):
        # 计算两颗恒星之间的引力产生的加速度
        dx = self.x - other_star.x
        dy = self.y - other_star.y
        distance_sq = dx ** 2 + dy ** 2

        if self.x < other_star.x:
            i_x = 1
        else:
            i_x = -1
        if self.y < other_star.y:
            i_y = 1
        else:
            i_y = -1

        # 防止两颗恒星重合导致除零错误
        if distance_sq != 0:
            force = G * other_star.mass / distance_sq
        else:
            print("碰撞检测到！")
            force = 0

        if dx != 0:  # 分母不为零
            ax_sq = force ** 2 / (((dy / dx) ** 2) + 1)
            self.ax += i_x * math.sqrt(ax_sq)
            self.ay += i_y * math.sqrt(force ** 2 - ax_sq)
        else:
            self.ay += force * i_y

    def run(self, time):
        # 根据加速度更新恒星的位置和速度
        self.vx += self.ax * time * 0.5
        self.vy += self.ay * time * 0.5
        self.x += self.vx * time
        self.y += self.vy * time

# 初始化
pygame.init()
screen = pygame.display.set_mode(WINDOW_SIZE, pygame.DOUBLEBUF, 32)
pygame.display.set_caption("三体问题模拟")
font = pygame.font.SysFont("simsunnsimsun", 16)

# 随机生成三颗恒星
stars = [
    Star(random.randint(100, 300), random.randint(100, 200), 0, 0, MASS, RADIUS),
    Star(random.randint(800, 1000), random.randint(100, 200), 0, 0, MASS, RADIUS),
    Star(random.randint(400, 600), random.randint(500, 600), 0, 0, MASS, RADIUS),
]

# 主循环
running = True
accuracy = INITIAL_ACCURACY  # 模拟精度

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # 根据精度进行多次迭代以提高模拟准确性
    for _ in range(accuracy):
        for i, star1 in enumerate(stars):
            star1.ax, star1.ay = 0, 0
            for j, star2 in enumerate(stars):
                if i != j:
                    star1.acceleration(star2)

        for star in stars:
            star.run(TIME_SLOT / accuracy)

    # 清屏
    screen.fill((0, 0, 0))
    # 计算最大速度，用于动态调整模拟精度
    max_v = max(math.sqrt(star.vx ** 2 + star.vy ** 2) for star in stars)

    # 绘制恒星和坐标信息
    for idx, star in enumerate(stars):
        color = (255, 100 * idx, 0)
        pygame.draw.circle(screen, color, (int(star.x), int(star.y)), star.radius)
        screen.blit(font.render(f"{int(star.x)},{int(star.y)}", True, THECOLORS["white"]), (0, idx * 15))

    # 动态调整精度
    accuracy = min(300, max(30, int(max_v))) * 5
    pygame.display.update()

sys.exit()